#pragma once
#include "TFGenerator.h"
#include <ANN/ANN.h>
#include "MatlabMatrix.h"
#include "MatlabDrawer.h"
using namespace matlab; 

#include <boost/thread/mutex.hpp>
#include "MetricLearner.h"
#include "OriginalTFGenerator.h"

/*
This class is quite similar to KNNGeneralizer. The difference is that this class maintains only one KNN 
regression for estimating the whole transition function, whereas the other one creates one 
regression problem for each output dimension. 

This file is taken from revision 126 of KNNGeneralizer. Therefore, if any bug is found for the other file, 
it should also be checked here!
*/

class KNNSimpleGeneralizer :
	public TFGenerator, MatlabDrawer
{
public:
	virtual ~KNNSimpleGeneralizer(void);

	KNNSimpleGeneralizer(MREAgent* p, int an, int od, taskspec_t& spec); 
	virtual void batchLearn( list<Transition>& history); 
	virtual void learn(const Transition* t); 
	virtual Observation predict(Observation st, Action a); 
	virtual Observation predict(Observation st, Action a, double& confidence); 
	virtual double getConfidence(Observation st, Action a);
	virtual void draw(); 

private:
	void resetTrees();				//free any storage used by the learners
	
	void updatePoints(std::list<Transition>& l, bool* use,  ANNpointArray& points, ANNpointArray& targets,  int action); //put transitions with action=action  from history l into points and targets. 
	
	int transformData(ANNpointArray& data,ANNpointArray& targets,  int size, int action); //transform data in data and target and return the new dimension
	
	void computeKernels(double* dists, double* kernels); 

	void drawScatteredPoints(int action); 
	void drawOption1(); 
	void drawOption2(); 

public:
	int K;											//how many nearest neighbors to use 
	double epsilon;									//epsilon for approximate KNN
	double sigma2;									//the sigma^2 used to compute kernel
	int DRAW_OPTION; 

	vector<ANNkd_tree*> learners;					//there is one learner for each action 
	vector<ANNpointArray> dataPoints;	    		//local copy of the points (it should not be changed between execution of batchlearn). 
	vector<ANNpointArray> targets;					//the labels for dataPoints (there's an array of observations for each action)
	int* reducedDimensions;							//an array of size action_number
	int* dataSize;									//an int array of size action_number that specifies how many data point for each action is available
	ANNidx* nnIdx;									//this is a temporary place to put the indexes of the k nearest neighbors
	ANNdist* dists;									//this is a temporary place to put the distances of the k nn

	boost::mutex m_mutex; 

	//warning: just for debugging, remove
	OriginalTFGenerator* otf; 

	MultiVariateMetricLearner* metricLearner;		
	friend class MetricLearner; 
};
